import websocket
import json
import ssl
import time
import threading
from datetime import datetime
import redis
from dotenv import load_dotenv
import os

pool = redis.ConnectionPool(host='localhost',port=6379, db=3)
r = redis.Redis(connection_pool=pool)

load_dotenv()
okx_get_orderbook_interval=os.getenv('OKX_ORDERBOOK_INTERVAL')
if okx_get_orderbook_interval:
    okx_get_orderbook_interval=int(okx_get_orderbook_interval)
else:
    print("环境变量 'okx_get_orderbook_interval' 未设置")

okx_token_str = os.getenv('OKX_TOKEN')
if okx_token_str:
    okx_token_array = okx_token_str.split(',')
    okx_token_array = [item.strip() for item in okx_token_array]
else:
    print("环境变量 'OKX_TOKEN' 未设置")

def on_error(ws, error):
    print(error)

def on_close(ws,close_status_code, close_msg):
    print("okx Connection closed")
    time.sleep(10)
    print("okx 正在重新连接")
    start_okx_websocket()

def on_open(ws,okx_token_array):
    print("okx websocket open")

def on_message(ws, message):
    global r
    if("pong"!=message):
        json_data=json.loads(message)
        if("action" in json_data and json_data["action"]=="snapshot"):
            trading_pair=json_data["arg"]["instId"]
            current_time = datetime.now().strftime('%Y-%m-%d-%H:%M:%S')
            trade_data = {
                "time": current_time,
                "trading_pair": trading_pair,
                "buy_order": json_data["data"][0]["bids"],
                "sell_order": json_data["data"][0]["asks"]
            }
            try:
                r.set(current_time+"-"+trading_pair+"-OKX-Orderbook", json.dumps(trade_data),ex=86400)
            except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
                print(f"Operation failed: {e}. Attempting to reconnect...")
                time.sleep(5)
                r = redis.Redis(connection_pool=pool)

        if("action" in json_data and json_data["action"]=="update"):
            grouping=json_data["arg"]["grouping"]
            trading_pair=json_data["arg"]["instId"]
            #取消订阅
            subscription_message = {
                "op":"unsubscribe",
                "args":[
                    {"channel":"books-grouped","instId":trading_pair,"grouping":grouping}
                ]
            }
            ws.send(json.dumps(subscription_message))
        # if("event" in json_data):
        #     print(message)


def send_message_periodically(ws, interval):
    while True:
        time.sleep(1)
        if ws.sock and ws.sock.connected:
            for token in okx_token_array:
                parts = token.rsplit('-', 1)
                instId = parts[0]  # "BTC-USDT-SWAP"
                grouping = parts[1]  # "0.1"
                subscription_message = {
                    "op":"subscribe",
                    "args":[
                        {"channel":"books-grouped","instId":instId,"grouping":grouping}
                    ]
                }
                ws.send(json.dumps(subscription_message))
                time.sleep(0.5)
        else:
            print("okx task 断开连接")
            break
        time.sleep(2)
        #ws.send(json.dumps(unsubscription_message))
        time.sleep(interval-3.5)
        


def send_ping(ws,interval):
    while True:
        time.sleep(interval)
        if ws.sock and ws.sock.connected:
            ws.send("ping")
        else:
            print("okx task 断开连接")
            break


def start_okx_websocket():
    websocket_url = "wss://wspri.okx.com:8443/ws/v5/ipublic"
    # 创建WebSocket应用程序对象
    try:
        ws = websocket.WebSocketApp(websocket_url,
                                    on_open=lambda ws: on_open(ws,okx_token_array=okx_token_array),
                                    on_message=on_message,
                                    on_error=on_error,
                                    on_close=on_close
                                    )
        ws_thread = threading.Thread(target=ws.run_forever,kwargs={"sslopt": {"cert_reqs": ssl.CERT_NONE}})
        ws_thread.start()

        time.sleep(3)
        if not ws.sock or not ws.sock.connected:
            raise ConnectionError("WebSocket connection failed")

        send_thread = threading.Thread(target=send_message_periodically, args=(ws, okx_get_orderbook_interval))
        send_thread.start()

        send_ping_thread = threading.Thread(target=send_ping, args=(ws, 20))
        send_ping_thread.start()

    except Exception as e:
        print(e)